[perf][rocm]: accelerate strict VIME rollout with direct paged CK and graph-safe execution - #390

Context

PR #388 established bitwise Qwen3/VIME train-rollout alignment on ROCm, but the strict R/R rollout path still carried substantial runtime overhead:

• paged KV was reconstructed into dense logical rows;
• GPU metadata was synchronized to the host in every decoder layer;
• Attention executed multiple Python-level row/KV-group dispatches;
• RoPE tables and weight layouts were repeatedly materialized;
• Attention and GEMM outputs passed through temporary buffers;
• deterministic collectives introduced staging copies and graph breaks.

PR #390 moves the strict ROCm rollout path from a materialization-heavy adapter to a direct-paged, graph-safe execution pipeline while preserving the train/rollout arithmetic contract.

Optimization architecture

flowchart TB
    INPUT(["VIME strict R/R rollout"]) --> META["Forward-scoped metadata cache"]

    subgraph HOTPATH["Strict ROCm compute hot path"]
        direction LR
        QKV["Deterministic QKV GEMM"] --> ROPE["Token-major deterministic RoPE"]
        ROPE --> ATTN["Direct paged AITER/CK"]
        ATTN --> OPROJ["Direct-output O-projection"]
    end

    META --> QKV
    OPROJ --> IPC["Stable-buffer IPC fixed-tree reduce"]
    IPC --> HEAD["Prepared LM-head and strict logprob"]
    HEAD --> GATE{"Exactness and provenance gate"}
    GATE --> OUTPUT(["Aligned rollout output"])

    subgraph CONTRACT["Strict arithmetic contract"]
        direction LR
        SPLIT["Split-KV disabled"] --> ACCUM["FP32 accumulation"]
        ACCUM --> TREE["Fixed reduction order"]
        TREE --> FALLBACK["No fallback"]
    end

    SPLIT -. "constrains" .-> ATTN
    TREE -. "constrains" .-> IPC
    FALLBACK -. "verified by" .-> GATE

    classDef runtime fill:#ddf4ff,stroke:#0969da,color:#24292f,stroke-width:2px
    classDef compute fill:#fff8c5,stroke:#bf8700,color:#24292f,stroke-width:2px
    classDef distributed fill:#dafbe1,stroke:#1a7f37,color:#24292f,stroke-width:2px
    classDef guard fill:#ffebe9,stroke:#cf222e,color:#24292f,stroke-width:2px

    class INPUT,META,OUTPUT runtime
    class QKV,ROPE,ATTN,OPROJ compute
    class IPC,HEAD distributed
    class SPLIT,ACCUM,TREE,FALLBACK,GATE guard

    style HOTPATH fill:#fffdf2,stroke:#bf8700,stroke-width:2px
    style CONTRACT fill:#fff5f5,stroke:#cf222e,stroke-width:2px

The solid arrows represent the steady-state rollout data path. The dashed arrows represent the strict arithmetic constraints applied to kernel dispatch, accumulation, communication, and runtime acceptance.

The optimized route keeps tensors on the GPU, preserves framework-owned storage, and removes redundant synchronization, layout conversion, temporary allocation, and output-copy operations from steady-state decode.

Key optimizations

Direct-paged Attention path

The optimized rollout path executes:

vLLM paged KV
    -> packed query metadata
    -> AITER mha_batch_prefill
    -> caller-owned output buffer

The direct path exposes the following runtime evidence:

paged_execution = direct_vllm_pages_to_aiter_batch_prefill_ck
dense_kv_materialized = false
query_schedule = paged_varlen_batch
core_launch_count = 1
split_kv = disabled
fallback = false

This removes dense KV reconstruction from the primary rollout path without introducing a moving Split-KV reduction schedule.

Exactness-preserving contract

The optimizations change data movement, dispatch, allocation, and storage placement. They do not relax the strict arithmetic contract:

• AITER/CK remains the production Attention arithmetic backend.
• Split-KV is disabled and num_splits is fixed at 1.
• Attention accumulation remains FP32.
• Output downcast occurs only at the final write.
• Dropout remains disabled.
• Deterministic GEMM retains its canonical reduction tree.
• Tensor-parallel collectives retain a fixed reduction order.
• Runtime fallback is forbidden.
• Backend identity, execution schedule, source fingerprint, and fallback state remain visible through provenance.
• Optimized results must remain byte-identical to the strict reference route.

30-step post-merge validation

The strict route was exercised for 30 matched P/P and R/R steps on the test integration line. Mismatch counts compare selected-token log probabilities between the training and rollout paths.

All 30 recorded steps are included in the averages below.

Train-rollout consistency

Across 1,920 comparisons, the strict R/R route remained byte-identical between training and rollout. The production P/P route produced 1,139 mismatches and a maximum absolute selected-token logprob difference of 1.041255.

Average performance across 30 steps

The 30-step result establishes the exactness property, but it does not establish performance parity with the production P/P route.

The remaining gap is concentrated in strict rollout execution, reference/logprob computation, and actor training. Weight synchronization is approximately at parity and is 1.27% faster in this measurement.

This result therefore serves both as end-to-end exactness evidence and as a baseline for the remaining ROCm optimization work.

Post-merge follow-up

PR #390 merged into test at fed8362.

The subsequent full-Graph work, route-specific AOT caching, position-plan optimization, and additional matched P/P versus R/R results are tracked in #393 so that results from different commit sets are not mixed into the isolated PR #390 performance claim.

The 30-step table above is retained as integration-level validation because it demonstrates that the strict route preserves zero mismatch under a longer matched workload while also making the remaining performance gap explicit.

Scope boundary

This PR optimizes the existing strict ROCm execution contract. It does not claim that every stage is faster than VIME production P/P, and it does not change arithmetic ordering solely for performance.

The current evidence establishes:

• direct-paged AITER/CK execution without dense KV materialization;
• fail-closed backend and schedule validation;
• byte-identical train-rollout selected-token log probabilities across 30 matched steps;
• a measurable reduction in weight-update time;
• a clearly quantified remaining performance gap against production P/P.

Potential follow-up work includes:

• profiling small-query mha_batch_prefill occupancy on gfx942;
• evaluating additional AITER paged kernels under a fixed non-Split-K contract;
• qualifying MFMA-specific deterministic GEMM schedules;
• reducing the remaining QKV, O-projection, LM-head, and logprob overhead;
• reducing strict actor-training overhead;
• publishing a final 200-step matched-workload comparison.

Related: #377, #385, #388, #393.